package com.hj.aiagent.core.agent;

import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.hj.aiagent.core.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * @author: hj
 * @description: 处理工具调用的基础代理类，具体实现了 think 和 act 方法，可以用作创建实例的父类
 * @date: 2025/7/16 14:14
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent {

    /**
     * 工具回调数组
     */
    private final ToolCallback[] toolCallbacks;

    /**
     * 保存工具调用信息的响应结果（要调用那些工具）
     */
    private ChatResponse toolCallChatResponse;

    /**
     * 工具调用管理者
     */
    private ToolCallingManager toolCallingManager;

    /**
     * 禁用 Spring AI 内置的工具调用机制，自己维护选项和消息上下文
     */
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] toolCallbacks) {
        super();
        this.toolCallbacks = toolCallbacks;
        // 禁用 Spring AI 内置的工具调用机制，自己维护选项和消息上下文
        this.chatOptions = DashScopeChatOptions.builder()
                .withProxyToolCalls(true)
                .build();
        this.toolCallingManager = ToolCallingManager.builder().build();
    }

    @Override
    public boolean think() {
        List<Message> messageList = getMemory();
        //1. 拼接提示词
        if (!getNextStepPrompt().isEmpty()) {
            messageList.add(new UserMessage(getNextStepPrompt()));
        }

        //2. 调用ai模型，并处理工具调用
        Prompt prompt = new Prompt(messageList, chatOptions);
        try {

            // 记录响应，用于等下 Act
            this.toolCallChatResponse = getClient().prompt(prompt).tools(toolCallbacks).call().chatResponse();
            // 3、解析工具调用结果，获取要调用的工具
            AssistantMessage assistantMessage = toolCallChatResponse.getResult().getOutput();
            //获取想要调用的工具
            List<AssistantMessage.ToolCall> toolCalls = assistantMessage.getToolCalls();
            //输出提示信息
            String result = assistantMessage.getText();
            log.info(getName() + "的思考：" + result);
            log.info(getName() + "选择了 " + toolCalls.size() + " 个工具来使用");
            String toolCallInfo = toolCalls.stream().map(toolCall -> String.format("%工具名称：%s，参数：%s", toolCall.name(), toolCall.arguments())).collect(Collectors.joining("\n"));
            log.info(toolCallInfo);

            //如果不需要工具，返回false
            if (toolCalls.isEmpty()) {
                // 只有不调用工具时，才需要手动记录助手消息
                messageList.add(assistantMessage);
                return false;
            } else {
                // 需要调用工具时，无需记录助手消息，因为调用工具时会自动记录
                return true;
            }
        } catch (Exception e) {
            log.error(getName() + "的思考过程遇到了问题：" + e.getMessage());
            messageList.add(new AssistantMessage("处理时遇到了错误：" + e.getMessage()));
            return false;
        }
    }

    @Override
    public String act() {
        if (!toolCallChatResponse.hasToolCalls()) {
            return "没有工具需要调用";
        }

        List<Message> messageList = getMemory();
        //调用工具
        Prompt prompt = new Prompt(messageList, chatOptions);
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallChatResponse);
        // 记录消息上下文，conversationHistory 已经包含了助手消息和工具调用返回的结果
        List<Message> conversationHistory = toolExecutionResult.conversationHistory();
        setMemory(conversationHistory);
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) conversationHistory.get(conversationHistory.size() - 1);
        //判断工具是否调用了终止工具
        boolean terminate = toolResponseMessage.getResponses()
                .stream()
                .anyMatch(response -> response.name().equals("doTerminate"));
        if (terminate){
            // 任务结束，更改状态
            setState(AgentState.FINISHED);
        }
        String results = toolResponseMessage.getResponses().stream()
                .map(response -> "工具 " + response.name() + " 返回的结果：" + response.responseData())
                .collect(Collectors.joining("\n"));
        log.info(results);
        return results;
    }
}
